{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1.import libs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"A script for data collection on Quafu platform.\"\"\"\n",
    "\n",
    "from utils.get_data import ShotTestReceiver, CircuitReceiver\n",
    "from qiskit import QuantumCircuit, ClassicalRegister\n",
    "from qiskit.qasm2 import dumps  # Use this for QASM conversion\n",
    "import itertools\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2.define some helpful functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_hadamard_test_circuit(qubit_indices: list) -> str:\n",
    "    \"\"\"Create a test circuit with Hadamard gates and CNOT gates.\n",
    "    \n",
    "    This function creates a quantum circuit that applies Hadamard gates and CNOT gates\n",
    "    to specified physical qubits and converts it to QASM format.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of physical qubit indices to map the logical qubits to\n",
    "        \n",
    "    Returns:\n",
    "        str: QASM representation of the circuit\n",
    "    \"\"\"\n",
    "    # Create circuit with enough qubits for the physical indices\n",
    "    max_qubit_idx = max(qubit_indices) + 1\n",
    "    circuit = QuantumCircuit(max_qubit_idx, len(qubit_indices))\n",
    "    \n",
    "    # Add gates using physical qubit indices directly\n",
    "    if len(qubit_indices) == 1:\n",
    "        circuit.h(qubit_indices[0])  # Single Hadamard gate\n",
    "    elif len(qubit_indices) >= 2:\n",
    "        # Apply Hadamard to first qubit followed by CNOTs\n",
    "        circuit.h(qubit_indices[0])  \n",
    "        for i in range(len(qubit_indices)-1):\n",
    "            circuit.cx(qubit_indices[i], qubit_indices[i+1])\n",
    "            \n",
    "    # Add measurements for all qubits\n",
    "    for i, qubit in enumerate(qubit_indices):\n",
    "        circuit.measure(qubit, i)\n",
    "    \n",
    "    return dumps(circuit)\n",
    "\n",
    "def create_random_circuit(qubit_indices: list, num_gates: int, valid_gates: list) -> str:\n",
    "    \"\"\"Create a random quantum circuit with specified gates (without parameter rotation gates).\n",
    "    \n",
    "    This function generates a random circuit by applying randomly selected gates\n",
    "    from the provided list to specified physical qubits.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of physical qubit indices to map the logical qubits to\n",
    "        num_gates: Number of gates to include in the circuit\n",
    "        valid_gates: List of valid quantum gates to use\n",
    "        \n",
    "    Returns:\n",
    "        str: QASM representation of the circuit\n",
    "    \"\"\"\n",
    "    max_qubit_idx = max(qubit_indices) + 1\n",
    "    circuit = QuantumCircuit(max_qubit_idx, len(qubit_indices))\n",
    "    \n",
    "    # Add random gates\n",
    "    for _ in range(num_gates):\n",
    "        gate = random.choice(valid_gates)\n",
    "        # Single qubit gates\n",
    "        if gate in ['x', 'y', 'z', 'h']:\n",
    "            qubits = random.sample(qubit_indices, 1)\n",
    "        # Two qubit gates\n",
    "        else:\n",
    "            qubits = random.sample(qubit_indices, 2)\n",
    "        circuit.append(gate, qubits)\n",
    "    \n",
    "    # Add measurements\n",
    "    for i, qubit in enumerate(qubit_indices):\n",
    "        circuit.measure(qubit, i)\n",
    "    \n",
    "    return dumps(circuit)\n",
    "\n",
    "def create_readout_benchmark_circuit(qubit_indices: list) -> list[str]:\n",
    "    \"\"\"Create benchmark circuits for readout error characterization.\n",
    "    \n",
    "    This function generates circuits that prepare all possible computational basis\n",
    "    states for the specified qubits to characterize readout errors.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of physical qubit indices to benchmark\n",
    "        \n",
    "    Returns:\n",
    "        list[str]: List of QASM strings for all state preparation combinations\n",
    "    \"\"\"\n",
    "    n_qubits = len(qubit_indices)\n",
    "    preparations = list(itertools.product([0,1], repeat=n_qubits))\n",
    "    circuits = []\n",
    "    \n",
    "    for prep in preparations:\n",
    "        max_qubit_idx = max(qubit_indices) + 1\n",
    "        circuit = QuantumCircuit(max_qubit_idx, n_qubits)\n",
    "        \n",
    "        # Prepare basis states\n",
    "        for qubit, state in zip(qubit_indices, prep):\n",
    "            if state == 1:\n",
    "                circuit.x(qubit)\n",
    "                \n",
    "        # Add measurements\n",
    "        for i, qubit in enumerate(qubit_indices):\n",
    "            circuit.measure(qubit, i)\n",
    "            \n",
    "        circuits.append(dumps(circuit))\n",
    "        \n",
    "    return circuits\n",
    "\n",
    "def create_readout_benchmark_circuit_with_qec(qubit_indices: list, \n",
    "                                            logical_qubit_indices: list) -> list[str]:\n",
    "    \"\"\"Create benchmark circuits with quantum error correction (QEC).\n",
    "    \n",
    "    This function generates circuits that implement error correction encoding\n",
    "    on top of basis state preparation for readout characterization.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of all physical qubit indices\n",
    "        logical_qubit_indices: List of logical qubit indices for encoding\n",
    "        \n",
    "    Returns:\n",
    "        list[str]: List of QASM strings for all encoded state preparations\n",
    "    \"\"\"\n",
    "    n_qubits = len(qubit_indices)\n",
    "    preparations = list(itertools.product([0,1], repeat=len(logical_qubit_indices)))\n",
    "    circuits = []\n",
    "    encoding_qubits = [i for i in qubit_indices if i not in logical_qubit_indices]\n",
    "\n",
    "    for prep in preparations:\n",
    "        max_qubit_idx = max(qubit_indices) + 1\n",
    "        circuit = QuantumCircuit(max_qubit_idx, n_qubits)\n",
    "\n",
    "        # Prepare and encode logical states\n",
    "        for qubit, state in zip(logical_qubit_indices, prep):\n",
    "            if state == 1:\n",
    "                circuit.x(qubit)\n",
    "            # Apply repetition code encoding\n",
    "            for i in encoding_qubits:\n",
    "                circuit.cx(qubit, i)\n",
    "\n",
    "        # Add measurements\n",
    "        for i, qubit in enumerate(qubit_indices):\n",
    "            circuit.measure(qubit, i)\n",
    "\n",
    "        circuits.append(dumps(circuit))\n",
    "\n",
    "    return circuits\n",
    "\n",
    "def create_repetition_code_circuit(qubit_indices: list) -> str:\n",
    "    \"\"\"Create a 3-qubit repetition code circuit.\n",
    "    \n",
    "    This function implements a simple 3-qubit repetition code with\n",
    "    a Hadamard gate on the first qubit followed by encoding operations.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of at least 3 physical qubit indices\n",
    "        \n",
    "    Returns:\n",
    "        str: QASM representation of the repetition code circuit\n",
    "    \"\"\"\n",
    "    if len(qubit_indices) < 3:\n",
    "        raise ValueError(\"Need at least 3 qubits for repetition code\")\n",
    "        \n",
    "    max_qubit_idx = max(qubit_indices) + 1\n",
    "    circuit = QuantumCircuit(max_qubit_idx, len(qubit_indices))\n",
    "    \n",
    "    # Initial Hadamard and encoding\n",
    "    circuit.h(qubit_indices[0])\n",
    "    circuit.cx(qubit_indices[0], qubit_indices[1])\n",
    "    circuit.cx(qubit_indices[0], qubit_indices[2])\n",
    "    \n",
    "    # Measurements\n",
    "    for i, qubit in enumerate(qubit_indices):\n",
    "        circuit.measure(qubit, i)\n",
    "        \n",
    "    return dumps(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.receive data of shot test  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "shot_test_platform_config = {\n",
    "    'platform': 'quafu',\n",
    "    'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'machine_name': 'Baihua',\n",
    "    'num_qubits': 2,\n",
    "    'exp_name': 'two-qubit_test_shots'\n",
    "}\n",
    "qubit_indices = [56, 57]\n",
    "shot_range = [int(1e3), 5*int(1e3), int(1e4), int(1e5)]\n",
    "shot_test_receiver = ShotTestReceiver(shot_test_platform_config, shot_range=shot_range, circuit_nums=1)\n",
    "shot_test_receiver.submit_experiment(create_hadamard_test_circuit(qubit_indices=qubit_indices))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "shot_test_receiver.receive_results()\n",
    "shot_test_receiver.save_results('results/20250304/one-qubit_h_test_shots/')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.receive data of benchmark test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_test_platform_config = {\n",
    "    'platform': 'quafu',\n",
    "    'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'machine_name': 'Baihua',\n",
    "    'num_qubits': 1,\n",
    "    'exp_name': 'one-qubit_benchmark'\n",
    "}\n",
    "\n",
    "benchmark_receiver = CircuitReceiver(benchmark_test_platform_config, circuit_nums=1, shots=100000, benchmark_flag=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_receiver.submit_experiment(create_readout_benchmark_circuit(qubit_indices=[56]))\n",
    "# benchmark_receiver.submit_experiment(create_readout_benchmark_circuit(qubit_indices=[56, 57]))\n",
    "# benchmark_receiver.submit_experiment(create_readout_benchmark_circuit(qubit_indices=[56, 55, 57]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_receiver.receive_results()\n",
    "benchmark_receiver.save_results('results/20250304/one-qubit_benchmark/')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.receive data of benchmark test with QEC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "qec_benchmark_platform_config = {\n",
    "    'platform': 'quafu',\n",
    "    'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'machine_name': 'Baihua',\n",
    "    'num_qubits': 3,\n",
    "    'exp_name': 'three-qubit_benchmark_qec'\n",
    "}\n",
    "\n",
    "qec_benchmark_receiver = CircuitReceiver(qec_benchmark_platform_config, circuit_nums=1, shots=100000, benchmark_flag=False)\n",
    "qec_benchmark_receiver.submit_experiment(create_readout_benchmark_circuit_with_qec(qubit_indices=[56, 55, 57], logical_qubit_indices=[56]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "qec_benchmark_receiver.receive_results()\n",
    "qec_benchmark_receiver.save_results('results/20250304/one-qubit_benchmark_qec/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6.receive data of repetition code test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "repetition_code_platform_config = {\n",
    "    'platform': 'quafu',\n",
    "    'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'machine_name': 'Baihua',\n",
    "    'num_qubits': 3,\n",
    "    'exp_name': 'three-qubit_repetition_code_hadamard'\n",
    "}\n",
    "\n",
    "repetition_code_receiver = CircuitReceiver(repetition_code_platform_config, circuit_nums=1, shots=100000, benchmark_flag=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['cx', 'cz', 'rx', 'ry', 'rz', 'x', 'y', 'z', 'h', 'delay', 'barrier']\n"
     ]
    }
   ],
   "source": [
    "repetition_code_receiver.submit_experiment(create_repetition_code_circuit(qubit_indices=[56, 55, 57]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "repetition_code_receiver.receive_results()\n",
    "repetition_code_receiver.save_results('results/20250304/three-qubit_repetition_code_hadamard/')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
